import lightgbm as lgb
from pro_data import DataSets
from utils.utiils_perf import *

def lgb_func(data_train, data_test, label_train):
    # 参数要设置小一些，防止产生过拟合
    param = {'num_leaves': 6,
             'objective': 'binary', 'min_data_in_leaf': 8,
             'metric': 'average_precision', 'max_depth': 3}
    num_round = 50

    # lgb
    data_train = lgb.Dataset(data=data_train,
                             label=label_train)

    bst = lgb.train(param,
                    data_train,
                    num_round)

    # test
    pred_test = bst.predict(data_test, num_iteration=bst.best_iteration)
    pred_test = [1 if i > 0.6 else 0 for i in pred_test]

    return pred_test

if __name__ == "__main__":
    path_x = "../data/data2_naca0012/dv.csv"
    path_y = "../data/data2_naca0012/fc.csv"

    ds = DataSets(path_x=path_x,
                  path_y=path_y,
                  label_index=2,
                  label_threshold=-0.16,
                  test_proportion=0.1)

    loop_cls(ds, lgb_func)
